/* Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * openUBMC is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *          http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

#ifdef G_ENABLE_TEST
#include <openssl/x509v3.h>
#include <openssl/pem.h>
#include <openssl/err.h>
#include <openssl/bio.h>
#endif // G_ENABLE_TEST
#include "spdm_command_impl.h"

// =================================================== SPDM公共函数 ===================================================
/*
 * @function：校验当前上下文是否合法
 */
bool CheckContext(void *spdm_context)
{
    return libspdm_check_context(spdm_context);
}

/*
 * @function：返回SPDM上下文的大小
 */
size_t GetContextSize(void)
{
    return libspdm_get_context_size();
}

/*
 * @function：返回SPDM上下文的大小，不包含安全消息上下文
 */
size_t GetContextSizeWithoutSecuredContext(void)
{
    return libspdm_get_context_size_without_secured_context();
}

/*
 * @function：初始化SPDM上下文
 */
uint32_t InitContext(void *spdm_context)
{
    return libspdm_init_context(spdm_context);
}

/*
 * @function：注册SPDM设备IO回调
 */
void RegisterDeviceIoFunc(void *spdm_context, libspdm_device_send_message_func send_message,
    libspdm_device_receive_message_func receive_message)
{
    libspdm_register_device_io_func(spdm_context, send_message, receive_message);
}

/*
 * @function：注册SPDM传输层编码/解码函数
 */
void RegisterTransportLayerFunc(
    void *spdm_context,
    uint32_t max_spdm_msg_size,
    uint32_t transport_header_size,
    uint32_t transport_tail_size,
    libspdm_transport_encode_message_func transport_encode_message,
    libspdm_transport_decode_message_func transport_decode_message)
{
    libspdm_register_transport_layer_func(spdm_context, max_spdm_msg_size, transport_header_size, transport_tail_size,
        transport_encode_message, transport_decode_message);
}

/*
 * @function：注册SPDM设备缓冲区管理函数
 */
void RegisterDeviceBufferFunc(
    void *spdm_context,
    uint32_t sender_buffer_size,
    uint32_t receiver_buffer_size,
    libspdm_device_acquire_sender_buffer_func acquire_sender_buffer,
    libspdm_device_release_sender_buffer_func release_sender_buffer,
    libspdm_device_acquire_receiver_buffer_func acquire_receiver_buffer,
    libspdm_device_release_receiver_buffer_func release_receiver_buffer)
{
    libspdm_register_device_buffer_func(spdm_context, sender_buffer_size, receiver_buffer_size, acquire_sender_buffer,
        release_sender_buffer, acquire_receiver_buffer, release_receiver_buffer);
}

/*
 * @function：设置SPDM上下文参数
 */
uint32_t SpdmSetData(void *spdm_context,
    libspdm_data_type_t data_type,
    const libspdm_data_parameter_t *parameter, const void *data,
    size_t data_size)
{
    return libspdm_set_data(spdm_context, data_type, parameter, data, data_size);
}

/*
 * @function：获取SPDM上下文参数
 */
uint32_t SpdmGetData(void *spdm_context, libspdm_data_type_t data_type,
    const libspdm_data_parameter_t *parameter, void *data, size_t *data_size)
{
    return libspdm_get_data(spdm_context, data_type, parameter, data, data_size);
}

/*
 * @function：清空指定内存
 */
void ZeroMem(void *buffer, size_t length)
{
    libspdm_zero_mem(buffer, length);
}

/*
 * @function：获取scratchbuffer的容量
 */
size_t GetRequiredScratchBufferSize(void *spdm_context)
{
    return libspdm_get_sizeof_required_scratch_buffer(spdm_context);
}

/*
 * @function：设置scratchbuffer地址和容量
 */
void SetScratchBuffer(void *spdm_context, void *scratch_buffer, size_t scratch_buffer_size)
{
    libspdm_set_scratch_buffer(spdm_context, scratch_buffer, scratch_buffer_size);
}

/*
 * @function：获取hash长度
 */
size_t GetHashSize(uint32_t base_hash_algo)
{
    return libspdm_get_hash_size(base_hash_algo);
}
// =================================================== SPDM请求端函数 ===================================================
/*
 * @function：初始化SPDM链接
 */
uint32_t InitSpdmConnection(void *spdm_context, bool get_version_only)
{
#ifdef G_ENABLE_TEST
    libspdm_data_parameter_t parameter = {};
    parameter.location = LIBSPDM_DATA_LOCATION_CONNECTION;

    uint32_t hash_algo = SPDM_ALGORITHMS_BASE_HASH_ALGO_TPM_ALG_SHA_256;
    SpdmSetData(spdm_context, LIBSPDM_DATA_BASE_HASH_ALGO, &parameter, &hash_algo, sizeof(uint32_t));

    uint16_t version = SPDM_MESSAGE_VERSION_10;
    SpdmSetData(spdm_context, LIBSPDM_DATA_SPDM_VERSION, &parameter, &version, sizeof(uint16_t));

    uint32_t asym_algo = SPDM_ALGORITHMS_BASE_ASYM_ALGO_TPM_ALG_RSAPSS_2048;
    SpdmSetData(spdm_context, LIBSPDM_DATA_BASE_ASYM_ALGO, &parameter, &asym_algo, sizeof(uint32_t));

    return LIBSPDM_STATUS_SUCCESS;
#else
    return libspdm_init_connection(spdm_context, get_version_only);
#endif
}

/*
 * @function：获取响应端设备上所有证书链的摘要
 */
uint32_t GetDigest(void *spdm_context, const uint32_t *session_id, uint8_t *slot_mask, void *total_digest_buffer)
{
#ifdef G_ENABLE_TEST
    *slot_mask = 0x7;
    memset_s(reinterpret_cast<uint8_t *>(total_digest_buffer), LIBSPDM_MAX_HASH_SIZE, 1, LIBSPDM_MAX_HASH_SIZE);
    memset_s(reinterpret_cast<uint8_t *>(total_digest_buffer) + LIBSPDM_MAX_HASH_SIZE, LIBSPDM_MAX_HASH_SIZE, 2,
        LIBSPDM_MAX_HASH_SIZE);
    memset_s(reinterpret_cast<uint8_t *>(total_digest_buffer) + 2 * LIBSPDM_MAX_HASH_SIZE, LIBSPDM_MAX_HASH_SIZE, 3,
        LIBSPDM_MAX_HASH_SIZE);
    return RET_OK;
#else
    return libspdm_get_digest(spdm_context, session_id, slot_mask, total_digest_buffer);
#endif // G_ENABLE_TEST
}

#ifdef G_ENABLE_TEST
static vector<uint8_t> get_mock_certificate()
{
    BIO *bio = BIO_new_mem_buf(test_cert_str.c_str(), -1);
    X509 *x509 = PEM_read_bio_X509(bio, nullptr, nullptr, nullptr);

    BIO *der_bio = BIO_new(BIO_s_mem());
    i2d_X509_bio(der_bio, x509);

    char *der_data;
    size_t der_length = BIO_get_mem_data(der_bio, &der_data);
    vector<uint8_t> buffer(der_data, der_data + der_length);

    // 在证书前扩展36个字节的header和hash信息
    vector<uint8_t> result(36 + buffer.size());
    fill(result.begin(), result.begin() + 36, 0x00);
    copy(buffer.begin(), buffer.end(), result.begin() + 36);

    BIO_free(bio);
    BIO_free(der_bio);
    X509_free(x509);

    return result;
}

#endif // G_ENABLE_TEST

/*
 * @function：获取响应端设备上指定槽位的证书链
 */
uint32_t GetCertificate(void *spdm_context, const uint32_t *session_id, uint8_t slot_id, size_t *cert_chain_size,
    void *cert_chain)
{
#ifdef G_ENABLE_TEST
    vector<uint8_t> result = get_mock_certificate();
    memcpy_s(cert_chain, LIBSPDM_MAX_CERT_CHAIN_SIZE, result.data(), result.size());
    *cert_chain_size = result.size();
    return RET_OK;
#else
    return libspdm_get_certificate(spdm_context, session_id, slot_id, cert_chain_size, cert_chain);
#endif // G_ENABLE_TEST
}

/*
 * @function：向响应端设备发送挑战以根据某个槽位中的密钥进行认证
 */
uint32_t ChallengeAuth(void *spdm_context, void *reserved, uint8_t slot_id, uint8_t measurement_hash_type,
    void *measurement_hash, uint8_t *slot_mask)
{
#ifdef G_ENABLE_TEST
    return LIBSPDM_STATUS_SUCCESS;
#else
    return libspdm_challenge(spdm_context, reserved, slot_id, measurement_hash_type, measurement_hash, slot_mask);
#endif // G_ENABLE_TEST
}

/*
 * @function：获取响应端设备的度量数据
 */
uint32_t GetMeasurement(void *spdm_context, const uint32_t *session_id, uint8_t request_attribute,
    uint8_t measurement_operation, uint8_t slot_id, uint8_t *content_changed, uint8_t *number_of_blocks,
    uint32_t *measurement_record_length, void *measurement_record)
{
#ifdef G_ENABLE_TEST
    if (measurement_operation == SPDM_GET_MEASUREMENTS_REQUEST_MEASUREMENT_OPERATION_TOTAL_NUMBER_OF_MEASUREMENTS) {
        *number_of_blocks = 3;
    } else if (measurement_operation == SPDM_GET_MEASUREMENTS_REQUEST_MEASUREMENT_OPERATION_ALL_MEASUREMENTS) {
        memcpy_s(measurement_record, DEFAULT_HASH_SIZE, g_all_measurement, DEFAULT_HASH_SIZE);
        *measurement_record_length = DEFAULT_HASH_SIZE;
    } else {
        auto it = g_measurements.find(measurement_operation);
        if (it == g_measurements.end()) {
            debug_log(DLOG_ERROR, "invalid operation %u", measurement_operation);
            return LIBSPDM_STATUS_INVALID_PARAMETER;
        }

        uint8_t mock_measurement[DEFAULT_HASH_SIZE] = {};
        copy(begin(it->second), end(it->second), mock_measurement);

        memcpy_s(measurement_record, DEFAULT_HASH_SIZE, mock_measurement, DEFAULT_HASH_SIZE);
        *measurement_record_length = DEFAULT_HASH_SIZE;
    }

    return LIBSPDM_STATUS_SUCCESS;
#else
    return libspdm_get_measurement(spdm_context, session_id, request_attribute, measurement_operation, slot_id,
        content_changed, number_of_blocks, measurement_record_length, measurement_record);
#endif
}

uint32_t GetMeasurementEx(void *spdm_context, const uint32_t *session_id,
    uint8_t request_attribute,
    uint8_t measurement_operation,
    uint8_t slot_id,
    uint8_t *content_changed,
    uint8_t *number_of_blocks,
    uint32_t *measurement_record_length,
    void *measurement_record,
    const void *requester_nonce_in,
    void *requester_nonce,
    void *responder_nonce,
    void *opaque_data,
    size_t *opaque_data_size)
{
#ifdef G_ENABLE_TEST
    memcpy_s(measurement_record, DEFAULT_HASH_SIZE, g_all_measurement, DEFAULT_HASH_SIZE);
    *measurement_record_length = DEFAULT_HASH_SIZE;

    return LIBSPDM_STATUS_SUCCESS;
#else
    return libspdm_get_measurement_ex(spdm_context, session_id, request_attribute, measurement_operation, slot_id,
        content_changed, number_of_blocks, measurement_record_length, measurement_record, requester_nonce_in,
        requester_nonce, responder_nonce, opaque_data, opaque_data_size);
#endif
}